package org.jiucheng.instrument;

import java.io.IOException;
import java.lang.instrument.ClassFileTransformer;
import java.lang.instrument.IllegalClassFormatException;
import java.security.ProtectionDomain;
import java.text.MessageFormat;

import org.jiucheng.agent.Agent;

import javassist.CannotCompileException;
import javassist.ClassPool;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.NotFoundException;

public class ExtClassFileTransformer implements ClassFileTransformer {
    
    private String[] classpaths;
    private String[] packages;
    private String exception;
    
    public ExtClassFileTransformer(String[] classpaths, String[] packages, String exception) {
        this.classpaths = classpaths;
        this.packages = packages;
        this.exception = exception;
    }
    
    public byte[] transform(ClassLoader loader, String className, Class<?> classBeingRedefined,
            ProtectionDomain protectionDomain, byte[] classfileBuffer) throws IllegalClassFormatException {
        String clazzName = className.replaceAll("/", ".");
        if(filterClassName(clazzName)) {
            try {
                ClassPool classPool = ClassPool.getDefault();
                for(String classpath : classpaths) {
                    classPool.appendClassPath(classpath);
                }
                CtClass ctClass = classPool.get(clazzName);
                CtClass e = classPool.get("java.lang.Exception");
                if(!ctClass.isInterface()) {
                    CtMethod[] ctMethods = ctClass.getDeclaredMethods();
                    for(CtMethod ctMethod : ctMethods) {
                        ctMethod.insertBefore(MessageFormat.format("org.jiucheng.instrument.Inject.before(\"{0}\", \"{1}\", $args);", ctClass.getName(), ctMethod.getName()));
                        if(clazzName.equals(exception)) {
                            ctMethod.addCatch(MessageFormat.format("org.jiucheng.instrument.Inject.exception(\"{0}\", \"{1}\", $e); throw $e;", ctClass.getName(), ctMethod.getName()), e);
                        }
                        ctMethod.insertAfter(MessageFormat.format("org.jiucheng.instrument.Inject.after(\"{0}\", \"{1}\", $args);", ctClass.getName(), ctMethod.getName()), true);
                        
                    }
                    ctClass.writeFile(Agent.getProp("webapp.dir") + "tmp");
                    return ctClass.toBytecode();
                }
            } catch (NotFoundException e) {
                e.printStackTrace();
            } catch (CannotCompileException e) {
                e.printStackTrace();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return classfileBuffer;
    }
    
    private boolean filterClassName(String className) {
        for(String pk : packages) {
            if(className.startsWith(pk)) {
                return true;
            }
        }
        return false;
    }
}
